import sys
import os
import numpy as np
import matplotlib.pyplot as plt
import codecs

log_dir = sys.argv[1]
def logFewShotPerf(log_dir):
    log_files = [fname for fname in os.listdir(log_dir) if ".log" in fname]
    for filename in log_files:
        log_file = os.path.join(log_dir, filename)
        fewshot_losses = []
        fewshot_accs = []
        tmp_losses, tmp_accs = [], []
        with codecs.open(log_file, 'r', encoding= u'utf-8',errors='ignore') as fr:
            for line in fr:
                if "-------> few shot data list ------>" in line:
                    if len(tmp_accs) != 0 and len(tmp_losses) != 0:
                        fewshot_losses.append(np.mean(tmp_losses))
                        fewshot_accs.append(np.mean(tmp_accs))
                    tmp_losses, tmp_accs = [], []
                if "####Few Shot" in line:
                    loss, acc = line.split("=")[1].split("/")
                    tmp_losses.append(float(loss))
                    tmp_accs.append(float(acc))
        if len(fewshot_accs) == 0 and len(fewshot_losses)==0:
            fewshot_accs, fewshot_losses = tmp_accs, tmp_losses
        steps = list(range(len(fewshot_accs)))
        plt.cla()
        plt.xlabel("steps")
        plt.plot(steps, fewshot_accs, label="accuracy")
        plt.legend()
        plt.savefig(f"{log_dir}/{filename.strip('.log')}_Acc.png")

        plt.cla()
        plt.xlabel("steps")
        plt.plot(steps, fewshot_losses, label="loss")
        plt.legend()
        plt.savefig(f"{log_dir}/{filename.strip('.log')}_Loss.png")

def plotContrastivePerf(log_file):
    distance = []
    with codecs.open(log_file, 'r', encoding=u'utf-8', errors='ignore') as fr:
        for line in fr:
            if "=======> Mutual Distance:" in line:
                dist = float(line.split("tensor(")[1].split(",")[0])
                distance.append(dist)

    steps = list(range(len(distance)))
    plt.xlabel("steps")
    plt.plot(steps, distance, label="mutual distance")
    plt.legend()
    plt.savefig(f"mutualDistance.png")

def plotDistillPerf(log_file):
    F1_list = []
    with codecs.open(log_file, 'r', encoding=u'utf-8', errors='ignore') as fr:
        for line in fr:
            if "Post-MetaTrain Performance of" in line:
                F1 = float(line.split(",")[6].strip("])"))
                F1_list.append(F1)
    model1_f1 = [f1 for idx, f1 in enumerate(F1_list) if idx%2!=1]
    model2_f1 = [f1 for idx, f1 in enumerate(F1_list) if idx%2==1]
    plt.xlabel("steps")
    plt.title("F1 Score")
    plt.plot(list(range(len(model1_f1))), model1_f1, label="model1")
    plt.plot(list(range(len(model2_f1))), model2_f1, label="model2")
    plt.legend()
    plt.savefig(f"F1Score.png")

# plotContrastivePerf(sys.argv[1])
plotDistillPerf(sys.argv[1])